library(data.table)
library(gridExtra)
library(grid)
library(tidyverse)
library(stringr)
library(latex2exp)

## Paper or presentation
paper <- TRUE

## data directory
dir.project <- './'
dir.data <- paste0(dir.project,'data/')
dir.raw <- paste0(dir.data,'raw/')
dir.generated <- paste0(dir.data,'generated/')

## figures directory for output
dir.fig.paper <- paste0(dir.project,'doc/paper/figures/decomposition/')
dir.fig.presentation <- paste0(dir.project,'results/figures/presentation/')

source(paste0(dir.project,'src/r/figures.R'))

fileName <- paste0(dir.project,'calibration_version.txt')
calibration_version <- read_file(fileName) %>% str_replace_all("[\r\n]" , "")

#############
## FIGURES ##
#############
# upper bound on price drop for plotting
pd.ub <- 70

dat.optimum <- fread(paste0(dir.generated,'matlab/OptimTax/gpops/','optimum_vars.csv')) %>% mutate(pricedrop = drop*100) %>% filter(pricedrop <= pd.ub)

formula_incentive_effects <- ggplot(dat.optimum , aes(x=pricedrop)) +
    geom_line(aes(y=I_M, linetype = '1')) +
    geom_line(aes(y=I_C, linetype = '2')) +
    theme_bw() +
    theme.serif +
    theme(legend.position = c(0.2,0.9), legend.background = element_rect(fill=alpha('white', 0))) + 
    xlab('pct. drop in robot price') +
    ylab('') +
    ylim(c(-2,2)) +
    scale_linetype_manual(name='',values=c('1'=1,'2'=2), labels=unname(TeX(c("$I_M$","$I_C$"))))

to.pdf(print(formula_incentive_effects),"formula_incentive_effects.pdf", width=width.double, height=height.double)

formula_effort_effects <- ggplot(dat.optimum , aes(x=pricedrop)) +
    geom_line(aes(y=xi_C_iM_sum, linetype = '1')) +
    geom_line(aes(y=xi_C_iC_sum, linetype = '2')) +
    theme_bw() +
    theme.serif +
    theme(legend.position = c(0.3,0.85), legend.background = element_rect(fill=alpha('white', 0))) + 
    xlab('pct. drop in robot price') +
    ylab('') +
    ylim(c(-0.2,0.2)) +
    scale_linetype_manual(name='',values=c('1'=1,'2'=2), labels=unname(TeX(c("$\\sum_{i} \\xi_{i}C_{iM}$","$\\sum_{i} \\xi_{i}C_{iC}$"))))

to.pdf(print(formula_effort_effects),"formula_effort_effects.pdf", width=width.double, height=height.double)

formula_elasticities <- ggplot(dat.optimum , aes(x=pricedrop)) +
    geom_line(aes(y=eps_Y_M_Y_R_K_B, linetype = '1')) +
    geom_line(aes(y=eps_Y_C_Y_R_K_B, linetype = '2')) +
    theme_bw() +
    theme.serif +
    theme(legend.position = c(0.3,0.9), legend.background = element_rect(fill=alpha('white', 0))) + 
    xlab('pct. drop in robot price') +
    ylab('') +
    ylim(c(0,0.7)) +
    scale_linetype_manual(name='',values=c('1'=1,'2'=2), labels=unname(TeX(c("$\\epsilon_{Y_M/Y_R,K_B}$","$\\epsilon_{Y_C/Y_R,K_B}$"))))

to.pdf(print(formula_elasticities),"formula_elasticities.pdf", width=width.double, height=height.double)

formula_part_effect <- ggplot(dat.optimum , aes(x=pricedrop)) +
    geom_line(aes(y=part_effect)) +
    theme_bw() +
    theme(legend.position = 'none') +
    theme.serif +
    xlab('pct. drop in robot price') +
    ylab('')

to.pdf(print(formula_part_effect),"formula_part_effect.pdf", width=width.double, height=height.double)

formula_tax_rev <- ggplot(dat.optimum , aes(x=pricedrop)) +
    geom_line(aes(y=robot_tax_rev)) +
    theme_bw() +
    theme(legend.position = 'none') +
    theme.serif +
    xlab('pct. drop in robot price') +
    ylab('')

to.pdf(print(formula_tax_rev),"formula_tax_rev.pdf", width=width.double, height=height.double)

formula_tax_base <- ggplot(dat.optimum , aes(x=pricedrop)) +
    geom_line(aes(y=robot_tax_base)) +
    theme_bw() +
    theme(legend.position = 'none') +
    theme.serif +
    xlab('pct. drop in robot price') +
    ylab('')

to.pdf(print(formula_tax_base),"formula_tax_base.pdf", width=width.double, height=height.double)

p <- ggplot(dat.optimum , aes(x=pricedrop)) +
    geom_line(aes(y=robot_tax_rev, linetype = '1')) +
    geom_line(aes(y=RHS_with_effects, linetype = '2')) +
    theme_bw() +
    theme(legend.position = 'none') +
    theme.serif +
    xlab('pct. drop in robot price') +
    ylab('')

p <- ggplot(dat.optimum , aes(x=pricedrop)) +
    geom_line(aes(y=tau_B, linetype = '1')) +
    geom_line(aes(y=RHS_with_effects/robot_tax_base, linetype = '2')) +
    theme_bw() +
    theme(legend.position = 'none') +
    theme.serif +
    xlab('pct. drop in robot price') +
    ylab('')

# Combined in one figure
plots_combined <- grid.arrange(formula_elasticities + ggtitle('A: Elasticities'),
                                   formula_incentive_effects + ggtitle('B: Incentive effects'),
                                   formula_effort_effects + ggtitle('C: Effort-reallocation effects'),
                                   formula_part_effect + ggtitle('D: Participation effects'),
                                   ncol = 2, nrow = 2)

ggsave(paste0(dir.fig,"formula_terms_combined.pdf"),plot=plots_combined,
       width=11,height=11,units="cm")

plots_combined <- grid.arrange(formula_tax_rev + ggtitle('A: Robot tax rev. in 1000 USD'),
                                   formula_tax_base + ggtitle('B: Robot tax base in 1000 USD'),
                                   ncol = 2, nrow = 1)

ggsave(paste0(dir.fig,"tax_rev_combined.pdf"),plot=plots_combined,
       width=13,height=5,units="cm")


## Decomposition

decomp <- dat.optimum %>% select(pricedrop)
decomp$tau_B_implied <- (dat.optimum$eps_Y_M_Y_R_K_B * (dat.optimum$I_M + dat.optimum$xi_C_iM_sum) + 
  dat.optimum$eps_Y_C_Y_R_K_B * (dat.optimum$I_C + dat.optimum$xi_C_iC_sum) +
  dat.optimum$part_residual) / (dat.optimum$q_B * dat.optimum$K_B) * 100

decomp$tau_B_only_incentive <- (dat.optimum$eps_Y_M_Y_R_K_B * (dat.optimum$I_M) + 
                                dat.optimum$eps_Y_C_Y_R_K_B * (dat.optimum$I_C)) / (dat.optimum$q_B * dat.optimum$K_B) * 100


decomp$tau_B_incentive_plus_effort <- (dat.optimum$eps_Y_M_Y_R_K_B * (dat.optimum$I_M + dat.optimum$xi_C_iM_sum) + 
                                       dat.optimum$eps_Y_C_Y_R_K_B * (dat.optimum$I_C + dat.optimum$xi_C_iC_sum)) / (dat.optimum$q_B * dat.optimum$K_B) * 100


p <- ggplot(decomp , aes(x=pricedrop)) +
    geom_line(aes(y=tau_B_only_incentive, linetype = 'I')) +
    geom_line(aes(y=tau_B_incentive_plus_effort, linetype = 'I+E')) +
    geom_line(aes(y=tau_B_implied, linetype = 'I+E+P')) +
    theme_bw() +
    theme(legend.position = 'none') +
    theme.serif +
    xlab('pct. drop in robot price') +
    ylab('') +
    scale_linetype_manual(name='',values=c('I'='dotted', 'I+E'='dashed', 'I+E+P'='solid'), labels=c("Incent. eff.", "Incent. + Effort-reall. eff.","Incent. + Effort-reall. + Part.  eff."))

to.pdf(
  fig.legend(p + theme(legend.position = 'top')),
  "legend_formula_decomposition.pdf",
  width = 5,
  height = 0.3
)

ggsave(paste0(dir.fig,"formula_decomposition.pdf"),plot=p,
       width=7,height=6,units="cm")
